package org.enhance.common.utils;

import java.util.*;
import java.util.concurrent.ConcurrentMap;

/**
 * 写时复制的map集合，来源于kafka源码
 *
 * @author wenpan 2024/8/31 10:38 AM
 */
public class CopyOnWriteMap<K, V> implements ConcurrentMap<K, V> {

    /**
     * 可以看到这里就是一个用 volatile 修饰的map，说白了CopyOnWriteMap就是对map进行了包装
     */
    private volatile Map<K, V> map;

    public CopyOnWriteMap() {
        map = Collections.emptyMap();
    }

    public CopyOnWriteMap(Map<K, V> map) {
        this.map = Collections.unmodifiableMap(map);
    }

    @Override
    public boolean containsKey(Object k) {
        return map.containsKey(k);
    }

    @Override
    public boolean containsValue(Object v) {
        return map.containsValue(v);
    }

    @Override
    public Set<java.util.Map.Entry<K, V>> entrySet() {
        return map.entrySet();
    }

    @Override
    public V get(Object k) {
        // 读操作，直接操作原map即可
        return map.get(k);
    }

    @Override
    public boolean isEmpty() {
        return map.isEmpty();
    }

    @Override
    public Set<K> keySet() {
        return map.keySet();
    }

    @Override
    public int size() {
        return map.size();
    }

    @Override
    public Collection<V> values() {
        return map.values();
    }

    @Override
    public synchronized void clear() {
        map = Collections.emptyMap();
    }

    @Override
    public synchronized V put(K k, V v) {
        // 往该map里添加数据的时候是将原来的数据拷贝一份
        Map<K, V> copy = new HashMap<>(map);
        // 将要添加的数据写到拷贝后的map里
        V prev = copy.put(k, v);
        // 用copy后的map替换原来的map
        map = Collections.unmodifiableMap(copy);
        return prev;
    }

    @Override
    public synchronized void putAll(Map<? extends K, ? extends V> entries) {
        // 先拷贝
        Map<K, V> copy = new HashMap<>(map);
        // 添加到拷贝的集合里
        copy.putAll(entries);
        // 替换为拷贝后的map
        map = Collections.unmodifiableMap(copy);
    }

    @Override
    public synchronized V remove(Object key) {
        // 移除动作，也是先拷贝一份
        Map<K, V> copy = new HashMap<>(map);
        // 拷贝后的结果进行移除
        V prev = copy.remove(key);
        // 替换
        map = Collections.unmodifiableMap(copy);
        return prev;
    }

    @Override
    public synchronized V putIfAbsent(K k, V v) {
        // 如果不存在说明要新写入了
        if (!containsKey(k)) {
            return put(k, v);
        } else {
            return get(k);
        }
    }

    @Override
    public synchronized boolean remove(Object k, Object v) {
        if (containsKey(k) && get(k).equals(v)) {
            remove(k);
            return true;
        } else {
            return false;
        }
    }

    @Override
    public synchronized boolean replace(K k, V original, V replacement) {
        if (containsKey(k) && get(k).equals(original)) {
            put(k, replacement);
            return true;
        } else {
            return false;
        }
    }

    @Override
    public synchronized V replace(K k, V v) {
        if (containsKey(k)) {
            return put(k, v);
        } else {
            return null;
        }
    }

}
